import os
import re
import logging
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import Tuple, List, Dict
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import r2_score, mean_squared_error, make_scorer
from sklearn.model_selection import cross_val_score
from sklearn.linear_model import LinearRegression
import matplotlib as mpl
from datetime import datetime
import matplotlib.font_manager as fm

from config import VISUALIZATION_CONFIG, MODEL_CONFIG

def setup_logging(log_file: str) -> None:
    """设置日志配置"""
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )

# 设置matplotlib中文字体
def setup_chinese_font():
    """设置中文字体支持"""
    try:
        # 直接使用字体名称列表
        plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Bitstream Vera Sans', 'Arial Unicode MS', 'WenQuanYi Micro Hei', 'WenQuanYi Zen Hei']
        plt.rcParams['font.family'] = 'sans-serif'
        plt.rcParams['axes.unicode_minus'] = False
        logging.info("设置中文字体成功")
        return True
    except Exception as e:
        logging.warning(f"设置中文字体失败: {e}")
        return False

def calculate_optimal_window_size(df: pd.DataFrame, ratio: float) -> int:
    """计算最优滑动窗口大小"""
    return max(int(len(df) * ratio), 5)  # 最小窗口大小为5

def calculate_optimal_threshold(df: pd.DataFrame, ratio: float) -> float:
    """计算最优速度差分阈值"""
    return df['Speed'].std() * ratio

def process_file(file_path: str, command: int) -> pd.DataFrame:
    """处理单个数据文件"""
    try:
        with open(file_path, 'r') as f:
            lines = f.readlines()

        extracted_data = []
        current_timestamp = None
        current_speed = None
        current_distance = None

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if line.startswith('Timestamp:'):
                current_timestamp = pd.to_datetime(line.split('Timestamp:')[1])
            elif line.startswith('Speed:'):
                speed_str = line.split('Speed:')[1].split(',')[0].strip()
                distance_str = line.split('Distance:')[1].strip()
                try:
                    current_speed = float(speed_str)
                    current_distance = float(distance_str)
                    if current_timestamp is not None:
                        extracted_data.append({
                            'Timestamp': current_timestamp,
                            'Speed': current_speed,
                            'Distance': current_distance
                        })
                        current_timestamp = None
                        current_speed = None
                        current_distance = None
                except ValueError as e:
                    logging.warning(f"无法解析数值: {line}, 错误: {e}")

        if not extracted_data:
            logging.warning(f"文件 {file_path} 没有有效数据")
            return pd.DataFrame()

        df = pd.DataFrame(extracted_data)
        df['Command'] = command
        logging.info(f"成功处理文件 {file_path}, 数据点数: {len(df)}")
        return df

    except Exception as e:
        logging.error(f"处理文件 {file_path} 时出错: {e}")
        return pd.DataFrame()

def find_steady_state(df: pd.DataFrame, window_size: int, diff_threshold: float) -> pd.DataFrame:
    """查找稳态数据"""
    # 计算速度差分
    df['Speed_Diff'] = df['Speed'].diff().abs()
    
    # 使用滚动窗口计算平均差分
    df['Avg_Abs_Diff'] = df['Speed_Diff'].rolling(window=window_size, center=True).mean()
    
    # 填充 NaN 值，避免数据丢失
    df['Avg_Abs_Diff'] = df['Avg_Abs_Diff'].fillna(df['Speed_Diff'].mean())
    
    # 如果阈值太小，可能导致没有稳态点，调整阈值
    if diff_threshold < 0.001:
        diff_threshold = 0.001
    
    # 寻找满足稳态条件的点
    steady_state_mask = df['Avg_Abs_Diff'] <= diff_threshold
    steady_state_df = df[steady_state_mask]
    
    # 如果稳态点太少，放宽条件
    if len(steady_state_df) < 10:
        # 尝试使用前25%的最小变化点作为稳态
        sorted_diff = df.sort_values(by='Avg_Abs_Diff')
        num_points = max(10, int(len(df) * 0.25))  # 至少10个点或25%的点
        steady_state_df = sorted_diff.head(num_points)
        logging.info(f"稳态点太少，已放宽条件选取前{num_points}个最小变化点")
    
    return steady_state_df

def plot_speed_data(df: pd.DataFrame, steady_state_df: pd.DataFrame, 
                   steady_state_speed: float, command: int, 
                   save_path: str) -> None:
    """绘制速度数据图"""
    # 设置matplotlib后端为Agg
    mpl.use('Agg')
    
    plt.figure(figsize=VISUALIZATION_CONFIG['figsize'], 
               dpi=VISUALIZATION_CONFIG['dpi'])
    
    # 转换时间戳为数值索引以避免多维索引问题
    time_indices = np.arange(len(df))
    
    # 获取当前时间作为时间戳
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 绘制原始速度曲线
    plt.plot(time_indices, df['Speed'].values, 'b-', label='Speed Data', alpha=0.5)
    
    # 如果有稳态数据，绘制稳态速度曲线和平均值
    if not steady_state_df.empty:
        # 找到稳态数据在原始数据中的索引
        steady_indices = []
        for idx, row in steady_state_df.iterrows():
            if idx in df.index:
                steady_indices.append(df.index.get_loc(idx))
        
        if steady_indices:
            # 绘制稳态速度点
            steady_x = [time_indices[i] for i in steady_indices]
            steady_y = [df.iloc[i]['Speed'] for i in steady_indices]
            plt.scatter(steady_x, steady_y, color='red', s=30, label='Steady State')
            
            # 绘制平均稳态速度线
            plt.axhline(y=steady_state_speed, color='g', linestyle='--', 
                       label=f'Avg:{steady_state_speed:.3f} m/s')
    
    # 添加简洁的标签和标题
    plt.xlabel('Sample Points')
    plt.ylabel('Speed (m/s)')
    plt.title(f'Speed Curve (Command={command})\n{current_time}')
    plt.legend()
    plt.grid(True)
    
    # 保存图片
    os.makedirs(save_path, exist_ok=True)
    plt.savefig(os.path.join(save_path, f'speed_plot_{command}.png'))
    plt.close()

def evaluate_model(model, X_train: np.ndarray, X_test: np.ndarray, 
                  y_train: np.ndarray, y_test: np.ndarray) -> Dict[str, float]:
    """评估模型性能"""
    y_pred_train = model.predict(X_train)
    y_pred_test = model.predict(X_test)
    
    return {
        'train_r2': r2_score(y_train, y_pred_train),
        'test_r2': r2_score(y_test, y_pred_test),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_pred_test))
    }

def analyze_model_complexity(X: np.ndarray, y: np.ndarray, 
                           max_degree: int = 5) -> Tuple[List[float], List[float]]:
    """分析模型复杂度"""
    train_scores = []
    test_scores = []
    
    for degree in range(1, max_degree + 1):
        poly = PolynomialFeatures(degree=degree)
        X_poly = poly.fit_transform(X)
        
        # 使用交叉验证评估模型
        cv_scores = cross_val_score(LinearRegression(), X_poly, y, 
                                  cv=MODEL_CONFIG['cv_folds'], 
                                  scoring=make_scorer(r2_score))
        
        train_scores.append(cv_scores.mean())
        test_scores.append(cv_scores.std())
    
    return train_scores, test_scores

def plot_combined_calibration_data(all_datasets, save_path: str) -> None:
    """绘制合并的标定数据图"""
    # 设置matplotlib后端为Agg
    mpl.use('Agg')
    
    plt.figure(figsize=VISUALIZATION_CONFIG['figsize'], 
               dpi=VISUALIZATION_CONFIG['dpi'])
    
    # 获取当前时间作为时间戳
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 创建曲线拟合对象和颜色映射
    colors = ['b', 'r', 'g', 'c', 'm', 'y', 'k']
    markers = ['o', 's', '^', 'D', 'v', '<', '>']
    
    # 绘制所有数据集
    combined_X = []
    combined_y = []
    
    for i, (dataset_name, dataset) in enumerate(all_datasets.items()):
        color_idx = i % len(colors)
        marker_idx = i % len(markers)
        
        X = dataset['data']['Command'].values
        y = dataset['data']['SteadyStateSpeed'].values
        
        # 绘制数据点
        plt.scatter(X, y, 
                    color=colors[color_idx], 
                    marker=markers[marker_idx],
                    s=50, 
                    alpha=0.7,
                    label=f"{dataset_name} Data")
        
        # 获取拟合曲线
        if 'model' in dataset and 'poly' in dataset:
            x_fit = np.linspace(min(X), max(X), 100).reshape(-1, 1)
            y_fit = dataset['model'].predict(dataset['poly'].transform(x_fit))
            plt.plot(x_fit, y_fit, 
                     color=colors[color_idx], 
                     linestyle='--',
                     alpha=0.7,
                     label=f"{dataset_name} Fit")
        
        # 收集所有数据点用于合并拟合
        combined_X.extend(X)
        combined_y.extend(y)
    
    # 如果有合并的模型，绘制合并拟合曲线
    if 'combined' in all_datasets and 'model' in all_datasets['combined'] and 'poly' in all_datasets['combined']:
        combined_model = all_datasets['combined']['model']
        combined_poly = all_datasets['combined']['poly']
        
        x_fit = np.linspace(min(combined_X), max(combined_X), 100).reshape(-1, 1)
        y_fit = combined_model.predict(combined_poly.transform(x_fit))
        
        plt.plot(x_fit, y_fit, 
                 color='black', 
                 linewidth=2,
                 label="Combined Fit")
    
    # 添加图表标题和标签
    plt.title(f"Calibration Results\n{current_time}")
    plt.xlabel("Command Value")
    plt.ylabel("Steady Speed (m/s)")
    plt.grid(True)
    plt.legend()
    
    # 保存图表
    plt.savefig(os.path.join(save_path, "combined_calibration_plot.png"))
    plt.close()
    logging.info(f"已保存合并数据标定结果图表")

def train_model(X, y):
    """训练多项式回归模型"""
    # 创建多项式特征
    degree = min(MODEL_CONFIG['degree'], len(X) - 1)  # 防止过拟合
    poly = PolynomialFeatures(degree=degree)
    X_poly = poly.fit_transform(X.reshape(-1, 1))
    
    # 训练模型
    model = LinearRegression()
    model.fit(X_poly, y)
    
    return model, poly 